"""
@Description :   工具
@Author      :   python_assignment_group 
@Time        :   2022/10/30 07:24:15
"""

import datetime
import os
import re
import shutil

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch

from configs import *


def flat_accuracy(preds, labels):
    """
    根据预测概率向量和标签计算准确率
    """
    pred_flat = np.argmax(preds, axis=1).flatten()
    labels_flat = labels.flatten()
    return np.sum(pred_flat == labels_flat) / len(labels_flat)


def format_time(elapsed):
    '''
    时间格式转换，用于将已用时间的格式设置为hh:mm:ss
    '''
    elapsed_rounded = int(round((elapsed)))
    return str(datetime.timedelta(seconds=elapsed_rounded))


def make_dir(path: str):
    """
    自定义新建文件夹，如果已存在就删除重建。
    """
    if os.path.exists(path):
        shutil.rmtree(path)
    os.makedirs(path)


def save_net_stats(path: str, net, optimizer, train_stat, now_epoch, mode=0):
    """
    网络状态存储
    """
    if not os.path.exists(path):
        os.makedirs(path)
    save_name = ""
    if mode == 0:
        save_name = os.path.join(path, "newest.pkl")
    elif mode == 1:  # 存储最优参数
        save_name = os.path.join(path, "best.pkl")

    # 创建断点
    check_point = {
        "net_stat_dicts": net.state_dict(),
        "optimizer_stat_dicts": optimizer.state_dict(),
        "train_stat": train_stat,
        "epoch": now_epoch
    }

    torch.save(check_point, save_name)


def load_net_stats(path: str, buffer_name: str, net, optimizer, mode="train"):
    """
    模型状态加载
    """
    buffer_name = os.path.join(path, buffer_name)
    if not os.path.exists(buffer_name):
        if mode == "train":
            print("未找到模型状态缓存，重新训练")
            return 0
        else:
            print("未找到模型状态缓存，请重新选择")
    check_point = torch.load(buffer_name)
    net.load_state_dict(check_point["net_stat_dicts"])
    if mode == "train":
        optimizer.load_state_dict(check_point["optimizer_stat_dicts"])
        min_epoch = check_point["epoch"]
        print("模型状态成功加载，从第{}个epoch继续训练".format(1+min_epoch))
        return min_epoch
    else:
        print("模型状态成功加载")
        return 0


def write_configs(file, config: dict, data_num: int):
    """
    把参数写入日志
    file:日志文件指针
    config:参数字典
    data:使用第几个数据集
    """
    assert file, "文件指针为空！"
    file.write("========config========\n")
    file.write("data_num:"+str(data_num)+"\n")
    for key, val in config.items():
        file.write(str(key)+":"+str(val)+"\n")


def plot_for_each_dataset(path: str, train_stat: dict):
    """
    train_stat = {
       "训练损失函数值"  : 每一个epoch中每一个batch存储一次
       "训练准确率"  : 每一个epoch中每40个batch存储一次
       "平均训练损失函数值" :  每一个epoch存储一次
       "平均训练准确率"  : 每一个epoch存储一次

       "验证损失函数值" : 每一个epoch中每一个batch存储一次
       "验证准确率" :  每一个epoch存储一次
       "平均验证损失函数值" : 每一个epoch存储一次

       "平均训练用时" :  每一个epoch存储一次
       "平均验证用时" :  每一个epoch存储一次
    }
    """
    """
    挑选以下内容进行画图(用时因设备性能不同,应该不太需要画吧)
    1.训练时
       1.1 训练损失函数值
       1.2 训练准确率
       1.3 平均训练损失函数值
       1.4 平均训练准确率
    2、验证时
       2.1 验证损失函数值
       2.2 验证准确率
       2.3 平均验证损失函数值
    """
    plt.style.use(["science", "grid", "no-latex", "cjk-sc-font"])
    plt.rcParams["font.sans-serif"] = ['KaiTi']
    plt.rcParams["axes.unicode_minus"] = False
    plt.rcParams["figure.figsize"] = (15.2, 8.8)
    matplotlib.use("Agg")
    if not os.path.exists(path):
        os.makedirs(path)

    # 训练时的四幅图
    plt.title("训练损失变化图")
    plt.plot(train_stat["训练损失函数值"])
    plt.xlabel('batch')
    plt.ylabel('损失函数值')
    plt.savefig(os.path.join(path, "训练损失变化图.pdf"))
    plt.clf()

    plt.title("训练准确率变化图")
    plt.plot(train_stat["训练准确率"])
    plt.xlabel('batch/40')
    plt.ylabel('准确率')
    plt.savefig(os.path.join(path, "训练准确率变化图.pdf"))
    plt.clf()

    plt.title("平均训练损失变化图")
    plt.plot(train_stat["平均训练损失函数值"])
    plt.xlabel('epoch')
    plt.ylabel('平均损失值')
    plt.savefig(os.path.join(path, "平均训练损失变化图.pdf"))
    plt.clf()

    plt.title("平均训练准确率变化图")
    plt.plot(train_stat["平均训练准确率"])
    plt.xlabel('epoch')
    plt.ylabel('平均准确率')
    plt.savefig(os.path.join(path, "平均训练准确率变化图.pdf"))
    plt.clf()

    # 验证时的两幅图
    plt.title("验证损失变化图")
    plt.plot(train_stat["验证损失函数值"])
    plt.xlabel('batch')
    plt.ylabel('损失函数值')
    plt.savefig(os.path.join(path, "验证损失变化图.pdf"))
    plt.clf()

    plt.title("验证准确率变化图")
    plt.plot(train_stat["验证准确率"])
    plt.xlabel('epoch')
    plt.ylabel('准确率')
    plt.savefig(os.path.join(path, "验证准确率变化图.pdf"))
    plt.clf()


def plot_for_each_net(path: str, net_train_stat: dict):
    # net_train_stat的键名为data1 data2 data3
    """
    将每一次数据集的 (都是一个epoch存储一次的信息)
    1、平均训练损失函数值
    2、平均训练准确率
    3、平均验证损失函数值
    4、验证准确率
    画出来
    """
    plt.style.use(["science", "grid", "no-latex", "cjk-sc-font"])
    plt.rcParams["font.sans-serif"] = ['KaiTi']
    plt.rcParams["axes.unicode_minus"] = False
    plt.rcParams["figure.figsize"] = (15.2, 8.8)
    matplotlib.use("Agg")
    if not os.path.exists(path):
        os.makedirs(path)

    plt.title('损失函数值变化图')
    plt.plot(net_train_stat['data1']['平均训练损失函数值'])
    plt.plot(net_train_stat['data1']['平均验证损失函数值'])
    plt.plot(net_train_stat['data2']['平均训练损失函数值'])
    plt.plot(net_train_stat['data2']['平均验证损失函数值'])
    plt.plot(net_train_stat['data3']['平均训练损失函数值'])
    plt.plot(net_train_stat['data3']['平均验证损失函数值'])
    plt.legend(['Data1_Train_Loss', 'Data1_Valid_Loss', 'Data2_Train_Loss',
               'Data2_Valid_Loss', 'Data3_Train_Loss', 'Data3_Valid_Loss'])
    plt.xlabel('epoch')
    plt.ylabel('损失函数值')
    plt.savefig(os.path.join(path, "损失函数值变化图.pdf"))
    plt.clf()

    plt.title('准确率变化图')
    plt.plot(net_train_stat['data1']['平均训练准确率'])
    plt.plot(net_train_stat['data1']['验证准确率'])
    plt.plot(net_train_stat['data2']['平均训练准确率'])
    plt.plot(net_train_stat['data2']['验证准确率'])
    plt.plot(net_train_stat['data3']['平均训练准确率'])
    plt.plot(net_train_stat['data3']['验证准确率'])
    plt.legend(['Data1_Train_Acc', 'Data1_Valid_Acc', 'Data2_Train_Acc',
               'Data2_Valid_Acc', 'Data3_Train_Acc', 'Data3_Valid_Acc'])
    plt.xlabel('epoch')
    plt.ylabel('准确率')
    plt.savefig(os.path.join(path, "准确率变化图.pdf"))
    plt.clf()


def plot_for_test(path: str, test_results: np.ndarray):
    """
    测试结果绘图
    """
    plt.style.use(["science", "grid", "no-latex", "cjk-sc-font"])
    plt.rcParams["font.sans-serif"] = ['KaiTi']
    plt.rcParams["axes.unicode_minus"] = False
    plt.rcParams["figure.figsize"] = (15.2, 8.8)
    sns.color_palette(palette="bright")
    matplotlib.use("Agg")
    if not os.path.exists(path):
        os.makedirs(path)
    test_results = pd.DataFrame(test_results.T, columns=nets_names)
    test_results.loc[3] = [test_results[nets_names[0]].mean(),test_results[nets_names[1]].mean(),test_results[nets_names[2]].mean()]
    test_results["数据集"] = ["Data1", "Data2", "Data3", "Average"]
    test_results = pd.melt(test_results, id_vars=["数据集"])
    test_results.columns = ["数据集", "网络", "准确率"]

    plt.title("网络测试准确率图")
    bar = sns.barplot(x="数据集", y="准确率", hue="网络",
                      data=test_results, palette="Greens")
    plt.bar_label(bar.containers[0])
    plt.bar_label(bar.containers[1])
    plt.bar_label(bar.containers[2])

    plt.savefig(os.path.join(path, "三个网络测试准确率图.pdf"))
    plt.clf()


def plot_from_saved_info(save_path=os.path.join("figs", "info_for_figs")):
    """
    根据存储的图像信息作图
    """
    files = os.listdir(save_path)
    for file in files:
        file_path = os.path.join(save_path, file)
        if file == "all_data.npy":
            continue
        elif file == "test_results.npy":
            test_results = np.load(file_path, allow_pickle=True)
            fig_path = os.path.join("figs", "test_result_figs")
            plot_for_test(fig_path, test_results)
        elif "all_data" in file:  # ...all_data.npy
            net_train_stat = np.load(file_path, allow_pickle=True).item()
            name = file.removesuffix("all_data.npy")
            name_ind = nets_names.index(name)
            fig_path = os.path.join(
                train_configs[name_ind]["figs_path"], "all_data")
            plot_for_each_net(fig_path, net_train_stat)
        else:  # ...data?.npy
            train_stat = np.load(file_path, allow_pickle=True).item()
            name = file.removesuffix(
                re.search(r"data[1-3]\.npy", file).group())
            name_ind = nets_names.index(name)
            fig_path = os.path.join(
                train_configs[name_ind]["figs_path"], "all_data")
            plot_for_each_dataset(fig_path, train_stat)
